from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
from numpy import arange

plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

param = np.array([
    [30, 50, 6, 18.09],
    [30, 50, 8, 17.04],
    [30, 50, 10, 19.45],
    [30, 70, 6, 9.11],
    [30, 70, 8, 9.4],
    [30, 70, 10, 8.85],
    [30, 90, 6, 9.35],
    [30, 90, 8, 9.4],
    [30, 90, 10, 9.08],
    [40, 50, 6, 30.33],
    [40, 50, 8, 31.2],
    [40, 50, 10, 29.2],
    [40, 70, 6, 24.12],
    [40, 70, 8, 24.34],
    [40, 70, 10, 23.11],
    [40, 90, 6, 36.01],
    [40, 90, 8, 35.35],
    [40, 90, 10, 35.66],
    [50, 50, 6, 6.41],
    [50, 50, 8, 6.62],
    [50, 50, 10, 6.73],
    [50, 70, 6, 6.41],
    [50, 70, 8, 6.62],
    [50, 70, 10, 6.73],
    [50, 90, 6, 8.38],
    [50, 90, 8, 8.3],
    [50, 90, 10, 8.29],
])


def f(T, H, SC, klist, khlist):
    k_1 = klist[0]
    k_3 = klist[1]
    k_4 = klist[2]
    k_5 = klist[3]
    k_6 = klist[4]

    S_C = khlist[0]
    S_S = khlist[1]
    k_H = khlist[2]
    alpha = khlist[3]
    # 变量
    solid_content = SC / 100  # 固含量
    T = T + 273.15  # 温度，单位为开尔文

    # 常量
    # 常量
    C_0 = 24e-3  # 任意的比例系数，用于确定DMF质量与时间的关系
    m_D_0 = 24e-3  # 初始时刻DMF的质量，单位为kg
    m_S = 6e-3  # 环丁砜总的质量，单位为g。无论是否溶解，环丁砜不会从体系中消失，总的质量是不变的
    m_C = (m_D_0 + m_S) * solid_content  # 初始时刻醋酸纤维素的质量，单位为kg，与初始固含量有关
    ro_D = 0.948e3  # DMF的密度，单位为kg/m^3
    ro_S = 1.261e3  # 环丁砜的密度，单位为kg/m^3
    ro_C = 1.3e3  # 醋酸纤维素的密度，单位为kg/m^3
    V = m_D_0 / ro_D + m_S / ro_S + m_C / ro_C  # 总体积，单位为m^3，认为不变
    A = 6.09451
    B = 2725.96
    C = 28.209  # 由温度计算饱和汽压时的三个常数
    A_0 = -k_1 * (10 ** (A - B / (T + C))) * 133.322 * (1 - k_H * H / 100) / ro_D / V
    n_correction_factor = 6  # 斯托克斯 - 爱因斯坦方程中的修正系数，对于大分子乃至宏观颗粒是6，对于小分子则体积越小，该值也会越小
    viscosity_rate = 0.92e-3  # 在20℃下DMF粘度与水粘度的比值
    k = 1.380649e-23  # 玻尔兹曼常数，单位为J/K
    r = 0.3413e-9  # 环丁砜的分子半径

    # 函数
    # DMF在蒸发的过程中，其质量随时间变化的函数
    def m_D(t):
        return C_0 * np.exp(A_0 * t)

    # 未能溶解而析出成液滴的环丁砜的质量随时间变化的函数
    def m_S_out(t):
        return max(0, m_S - m_D(t) * S_S)

    # 未能溶解而析出成固体的醋酸纤维素的质量随时间变化的函数
    def m_C_out(t):
        return max(0, m_C - m_D(t) * S_C)

    # 析出的醋酸纤维素的体积占比随时间变化的函数
    def phi_C_out(t):
        return m_C_out(t) / ro_C / V

    # 由水在30、40、50（℃）下的粘度和DMF与水粘度的关系计算不同温度下DMF的粘度，单位为mPa*s
    def eta_0(temp):
        if temp == 30 + 273.15:
            return 0.8007 * viscosity_rate
        elif temp == 40 + 273.15:
            return 0.6560 * viscosity_rate
        elif temp == 50 + 273.15:
            return 0.5494 * viscosity_rate

    # 不同温度下混合溶液总的粘度随时间变化的函数
    def eta(temp, t):
        return eta_0(temp) * (1 + 2.5 * phi_C_out(t))

    # 扩散系数。假定一个小液滴内仅有一个分子，且分子为球形
    def D(temp, t):
        return k * temp / (n_correction_factor * np.pi * r * eta(temp, t))

    # 液滴运动的平均速度
    def v(temp, t):
        return k_3 * (D(temp, t) ** 0.5)

    # 析出的环丁砜的分子数密度随时间变化的函数
    def n_molecular_number_density(t):
        return m_S_out(t) / (ro_S * 4 / 3 * np.pi * r ** 3) / V

    # 小液滴间的碰撞频率
    def Z(temp, t):
        return (2 ** 0.5) * n_molecular_number_density(t) * np.pi * (r ** 2) * v(temp, t)

    tf = 0
    for t in np.arange(0, 500, 0.01):
        if m_D(t) < 0.024 * alpha:
            tf = t
            break

    dy = lambda m_s, t: k_4 * Z(T, t) + k_5 * v(T, t) * (m_S_out(t) - m_s) / V * m_s
    t = arange(1, tf, 0.01)
    sol = odeint(dy, 0, t)

    if len(sol.T[0]) == 0:
        re = 10000
    else:
        re = k_6 * sol.T[0][-1]
    return re


def pred(klist, khlist):
    predictions = []
    for i in range(param.shape[0]):
        T = param[i][0]
        H = param[i][1]
        SC = param[i][2]
        re = f(T, H, SC, klist, khlist)
        predictions.append(re)
    return predictions


def target():
    targets = []
    for i in range(param.shape[0]):
        t = param[i][-1]
        targets.append(t)
    return targets


def loss(pred, target):
    loss = 0
    for i in range(param.shape[0]):
        loss += np.linalg.norm(pred[i] - target[i])
    return loss


def printre(klist, khlist):
    t = target()
    for i in range(param.shape[0]):
        T = param[i][0]
        H = param[i][1]
        SC = param[i][2]
        re = f(T, H, SC, klist, khlist)
        tt = t[i]
        print(f"预测值{re},真实值{tt}")


def plot_SA_history():
    plt.figure(figsize=(14, 10))

    # 子图1：损失函数变化
    plt.subplot(2, 2, 1)
    plt.semilogy(history['best_loss'], 'r-', label='Best Loss')
    plt.semilogy(history['current_loss'], 'b--', alpha=0.5, label='Current Loss')
    plt.xlabel('Iteration')
    plt.ylabel('Loss')
    plt.title('Loss变化曲线')
    plt.legend()
    plt.grid(True, which="both", ls="--")

    # 子图2：温度衰减曲线
    plt.subplot(2, 2, 2)
    plt.plot(history['temperature'], 'g-')
    plt.xlabel('Iteration')
    plt.ylabel('Temperature')
    plt.title('温度下降曲线')
    plt.grid(True, ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 2, 3)
    k_array = np.array(history['k_params'])
    for i in range(5):
        plt.plot(k_array[:, i], label=f'k_{i + 1}')
    plt.xlabel('Iteration')
    plt.ylabel('k参数值')
    plt.title('k参数变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    # 子图4：kh参数演化
    plt.subplot(2, 2, 4)
    kh_array = np.array(history['kh_params'])
    for i in range(3):
        plt.plot(kh_array[:, i], label=f'kh_{i + 1}')
    plt.xlabel('Iteration')
    plt.ylabel('物理参数值')
    plt.title('物理参数变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    plt.tight_layout()
    plt.show()


# 在SA函数外部初始化记录容器
history = {
    'temperature': [],
    'best_loss': [],
    'current_loss': [],
    'k_params': [],
    'kh_params': []
}


k_bound = np.array([
    [0, 1],
    [0.001, 2],
    [0.001, 1],
    [0.001, 1],
    [0.001, 1],
])


kh_bound = np.array([
    [0.5, 2],
    [0.5, 1],
    [0, 0.8],
    [0, 1]
])

def SA(t0, tf, alpha, iter):
    global history
    t = t0
    k0 = np.array([0.02, 1, 0.01, 0.1, 0.5])
    kh0 = np.array([0.8, 0.9, 0.5, 0.6])
    '''
    lb190.91573602984445
    [0.01905417 1.37293927 0.00612436 0.0921701  0.87936372]
    [0.55282728 0.93158048 0.         0.12485199]
    '''
    kc = k0
    khc = kh0
    lc = loss(pred(kc, khc), target())

    kb = k0
    khb = kh0
    lb = loss(pred(kc, khc), target())
    for i in range(iter):
        kn = [j for j in kc]
        kn[0] += np.random.normal(0, 0.001)
        kn[1] += np.random.normal(0, 0.1)
        kn[2] += np.random.normal(0, 0.001)
        kn[3] += np.random.normal(0, 0.01)
        kn[4] += np.random.normal(0, 0.1)
        # kn = kc + np.random.normal(0, 0.001, size=5)
        kn = np.clip(kn, k_bound[:, 0], k_bound[:, 1])
        # khn = khc + np.random.uniform(-0.1, 0.1, size=3)
        # khn = np.clip(khn, kh_bound[:, 0], kh_bound[:, 1])
        khn = [j for j in khc]
        khn[0] += np.random.uniform(-0.1, 0.1)
        khn[1] += np.random.uniform(-0.1, 0.1)
        khn[2] += np.random.uniform(-0.1, 0.1)
        khn[3] += np.random.uniform(-0.1, 0.1)
        khn = np.clip(khn, kh_bound[:, 0], kh_bound[:, 1])
        pre = pred(kn, khn)
        ln = loss(pre, target())
        print(f"kn:{kn},khn:{khn}")
        if ln < lc or np.random.rand() < np.exp(-(ln - lc) / t):
            kc = kn
            khc = khn
            lc = ln
            if lc < lb:
                kb = kc
                khb = khn
                lb = lc
        t *= alpha
        if t < tf:
            break
        print(f"iter{i}, lb{lb},kc{kc},khc{khc},lc{lc},")
        history['temperature'].append(t)
        history['best_loss'].append(lb)
        history['current_loss'].append(lc)
        history['k_params'].append(kc.copy())
        history['kh_params'].append(khc.copy())
    return kb, khb


re1, re2 = SA(800, 0.001, 0.95, 1000)
print(re1, re2)
plot_SA_history()
printre(re1, re2)
